import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


# 定义目标函数：f(x, y) = (x - 1)^2 + (y - 2)^2
def objective(x, y):
    return (x - 4) ** 2 + (y + 2) ** 2


# 定义梯度函数
def gradient(x, y):
    df_dx = 2 * (x - 1)
    df_dy = 2 * (y - 2)
    return df_dx, df_dy


# 梯度下降算法
def gradient_descent(start_x, start_y, learning_rate, n_iterations):
    path = []
    x, y = start_x, start_y
    path.append((x, y))

    for i in range(n_iterations):
        grad_x, grad_y = gradient(x, y)
        x -= learning_rate * grad_x
        y -= learning_rate * grad_y
        path.append((x, y))
        print('第 %d 轮迭代，当前位置: (%.4f, %.4f)' % (i + 1, x, y))

    return path


# 初始值和超参数
start_x, start_y = 0, 0
learning_rate = 0.1
n_iterations = 50

# 运行梯度下降
path = gradient_descent(start_x, start_y, learning_rate, n_iterations)

# 准备绘图数据
x_vals = np.linspace(-5, 7, 100)
y_vals = np.linspace(-5, 7, 100)
X, Y = np.meshgrid(x_vals, y_vals)
Z = objective(X, Y)

# 将路径转换为 NumPy 数组以便绘制
path = np.array(path)
Z_path = objective(path[:, 0], path[:, 1])

# 创建 3D 绘图
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')

# 绘制目标函数的曲面
ax.plot_surface(X, Y, Z, cmap='viridis', alpha=0.6)

# 绘制梯度下降路径上的所有点
for i in range(len(path)):
    ax.scatter(path[i, 0], path[i, 1], Z_path[i], color=(i / len(path), 0, 1 - i / len(path)), s=50)
    # 可选：绘制从当前点指向最小值的箭头（帮助理解方向）
    ax.plot([path[i, 0], 1], [path[i, 1], 2], [Z_path[i], 0], color='gray', linestyle='--', linewidth=0.8)

# 设置初始点
ax.scatter(start_x, start_y, objective(start_x, start_y), color='blue', s=100, label='起点')

# 设置最小值点
ax.scatter(1, 2, 0, color='green', s=100, label='最小值点 (1, 2)')

# 设置标签和标题
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('f(x, y)')
ax.set_title('二元函数的梯度下降过程（每步梯度点）')

# 显示图例
ax.legend()

# 显示图形
plt.show()
